from mcp.server.fastmcp import FastMCP, Context
import subprocess
import json
from typing import Dict, Any

mcp = FastMCP("Kubernetes MCP Server")

"""
Kubernetes MCP Server
面向 Kubernetes 命令行操作

MCP tools 列表：
1. kube_apiserver_service_management - 管理 kube-apiserver 服务状态
   参数：action (str) - 操作类型：restart, enable, reload, check_logs
   返回：统一的 JSON 结构

2. kube_controller_manager_service_management - 管理 kube-controller-manager 服务状态
   参数：action (str) - 操作类型：restart, enable, reload, check_logs
   返回：统一的 JSON 结构

3. kube_proxy_service_management - 管理 kube-proxy 服务状态
   参数：action (str) - 操作类型：execution, reload
   返回：统一的 JSON 结构

4. kube_scheduler_service_management - 管理 kube-scheduler 服务状态
   参数：action (str) - 操作类型：execution, reload
   返回：统一的 JSON 结构

5. kubelet_service_management - 管理 kubelet 服务状态
   参数：action (str) - 操作类型：restart, enable, reload, check_logs
   返回：统一的 JSON 结构
"""

def run_command(command: str) -> Dict[str, Any]:
    """执行 shell 命令并返回统一格式的结果"""
    try:
        result = subprocess.run(
            command,
            shell=True,
            capture_output=True,
            text=True,
            timeout=30
        )
        return {
            "success": result.returncode == 0,
            "command": command,
            "exit_code": result.returncode,
            "stdout": result.stdout,
            "stderr": result.stderr
        }
    except subprocess.TimeoutExpired:
        return {
            "success": False,
            "command": command,
            "exit_code": -1,
            "stdout": "",
            "stderr": "Command timed out"
        }
    except Exception as e:
        return {
            "success": False,
            "command": command,
            "exit_code": -1,
            "stdout": "",
            "stderr": str(e)
        }

@mcp.tool()
def kube_apiserver_service_management(action: str) -> Dict[str, Any]:
    """
    管理 kube-apiserver 服务状态
    
    参数:
        action: 操作类型 - restart(重启服务), enable(启用服务), reload(重载服务), check_logs(检查日志错误)
    
    返回:
        JSON 结构: {
            "success": bool,
            "command": str,
            "exit_code": int,
            "stdout": str,
            "stderr": str
        }
    """
    service = "kube-apiserver.service"
    
    if action == "restart":
        return run_command(f"systemctl restart {service}")
    elif action == "enable":
        return run_command(f"systemctl enable {service}")
    elif action == "reload":
        return run_command(f"systemctl reload {service}")
    elif action == "check_logs":
        # 检查日志中的错误信息，过滤掉调试和信息级别的日志
        return run_command(f"journalctl -u {service} | grep -i \"fail\\|error\" | grep -v -i \"DEBUG\\|INFO\\|WARNING\" | grep -v \"AdditionalErrorMsg\"")
    else:
        return {
            "success": False,
            "command": f"Unknown action: {action}",
            "exit_code": -1,
            "stdout": "",
            "stderr": f"Unsupported action: {action}. Supported actions: restart, enable, reload, check_logs"
        }

@mcp.tool()
def kube_controller_manager_service_management(action: str) -> Dict[str, Any]:
    """
    管理 kube-controller-manager 服务状态
    
    参数:
        action: 操作类型 - restart(重启服务), enable(启用服务), reload(重载服务), check_logs(检查日志错误)
    
    返回:
        JSON 结构: {
            "success": bool,
            "command": str,
            "exit_code": int,
            "stdout": str,
            "stderr": str
        }
    """
    service = "kube-controller-manager.service"
    
    if action == "restart":
        return run_command(f"systemctl restart {service}")
    elif action == "enable":
        return run_command(f"systemctl enable {service}")
    elif action == "reload":
        return run_command(f"systemctl reload {service}")
    elif action == "check_logs":
        # 检查日志中的错误信息，过滤特定错误
        return run_command(f"journalctl -u {service} | grep -i \"fail\\|error\" | grep -v -i \"DEBUG\\|INFO\\|WARNING\" | grep -Ev \"authentication-tolerate-lookup-failure|start cloud node lifecycle controller\"")
    else:
        return {
            "success": False,
            "command": f"Unknown action: {action}",
            "exit_code": -1,
            "stdout": "",
            "stderr": f"Unsupported action: {action}. Supported actions: restart, enable, reload, check_logs"
        }

@mcp.tool()
def kube_proxy_service_management(action: str) -> Dict[str, Any]:
    """
    管理 kube-proxy 服务状态
    
    参数:
        action: 操作类型 - execution(执行服务), reload(重载服务)
    
    返回:
        JSON 结构: {
            "success": bool,
            "command": str,
            "exit_code": int,
            "stdout": str,
            "stderr": str
        }
    """
    service = "kube-proxy.service"
    
    if action == "execution":
        return run_command(f"systemctl start {service} && systemctl status {service}")
    elif action == "reload":
        return run_command(f"systemctl reload {service}")
    else:
        return {
            "success": False,
            "command": f"Unknown action: {action}",
            "exit_code": -1,
            "stdout": "",
            "stderr": f"Unsupported action: {action}. Supported actions: execution, reload"
        }

@mcp.tool()
def kube_scheduler_service_management(action: str) -> Dict[str, Any]:
    """
    管理 kube-scheduler 服务状态
    
    参数:
        action: 操作类型 - execution(执行服务), reload(重载服务)
    
    返回:
        JSON 结构: {
            "success": bool,
            "command": str,
            "exit_code": int,
            "stdout": str,
            "stderr": str
        }
    """
    service = "kube-scheduler.service"
    
    if action == "execution":
        return run_command(f"systemctl start {service} && systemctl status {service}")
    elif action == "reload":
        return run_command(f"systemctl reload {service}")
    else:
        return {
            "success": False,
            "command": f"Unknown action: {action}",
            "exit_code": -1,
            "stdout": "",
            "stderr": f"Unsupported action: {action}. Supported actions: execution, reload"
        }

@mcp.tool()
def kubelet_service_management(action: str) -> Dict[str, Any]:
    """
    管理 kubelet 服务状态
    
    参数:
        action: 操作类型 - restart(重启服务), enable(启用服务), reload(重载服务), check_logs(检查日志错误)
    
    返回:
        JSON 结构: {
            "success": bool,
            "command": str,
            "exit_code": int,
            "stdout": str,
            "stderr": str
        }
    """
    service = "kubelet.service"
    
    if action == "restart":
        return run_command(f"systemctl restart {service}")
    elif action == "enable":
        return run_command(f"systemctl enable {service}")
    elif action == "reload":
        return run_command(f"systemctl reload {service}")
    elif action == "check_logs":
        # 检查日志中的错误信息，过滤特定错误
        return run_command(f"journalctl -u {service} | grep -i \"fail\\|error\" | grep -v -i \"DEBUG\\|INFO\\|WARNING\" | grep -vE \"CredentialsChainVerboseErrors|Failed to read data from checkpoint|Image garbage collection failed once\"")
    else:
        return {
            "success": False,
            "command": f"Unknown action: {action}",
            "exit_code": -1,
            "stdout": "",
            "stderr": f"Unsupported action: {action}. Supported actions: restart, enable, reload, check_logs"
        }

if __name__ == "__main__":
    mcp.run()